#include <cstdlib>
#include <cmath>
#include <iostream>

#include "ConfigFile.h"

#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/generate.h>
#include <thrust/functional.h>
#include <thrust/random.h>
#include <thrust/transform.h>
#include <thrust/sort.h>
#include <thrust/reduce.h>
#include <thrust/iterator/constant_iterator.h>
#include <thrust/scatter.h>

#include "util.h"
#include "simulation.h"

// Functors and other items for transformations and generating values

// generates a random float uniformly distributed in [-1.0,1.0)
float random_float(void)
{
  static thrust::default_random_engine rng;
  thrust::uniform_real_distribution<float> u(-1.0, 1.0);

  return u(rng);
}


// functor returns a perturbed version of the input based on
// normal distribution with mean 0, std_dev as defined in initialization
// only works on host vectors
struct perturb_with_normal : thrust::unary_function<float, float>
{
  perturb_with_normal(const float std_dev) : sd(std_dev)
  {}

//  __host__ __device__
  float operator()(float x) const
  {
    static thrust::minstd_rand rng;
    thrust::random::experimental::normal_distribution<float> dist(0.0f, sd);
    
    return dist(rng);
  }

  float sd;
};


// command line version info
void print_version(void)
{
  std::cout << "--------------------------------------------------------------------------------" << std::endl;
  std::cout << "CUDA Creaturs Version 0.04" << std::endl;
  std::cout << "By: Andrew Hershberger & Vadim Ogievetsky" << std::endl;
  std::cout << "Stanford CS / 19 May 2010" << std::endl;
}

// command line help info
void print_help(void)
{
  std::cout << "--------------------------------------------------------------------------------" << std::endl;
  std::cout << "Command line parameters:" << std::endl;
  std::cout << "  --help          -h : Prints this help and exits" << std::endl;
  std::cout << "  --version          : Prints the version information" << std::endl;
  std::cout << "  --config-in     -ci: The input configuration filename" << std::endl;
  std::cout << "  --config-out    -co: The saved configuration filename" << std::endl;
  std::cout << "  --verbose       -v : Sets mode to verbose" << std::endl;
  std::cout << "  --cpu              : Runs the sim on the CPU to check results" << std::endl;
  std::cout << "  --rounds        -r : The number of rounds to run per itteration" << std::endl;
  std::cout << "  --memory        -m : The amount of memory each creature has" << std::endl;
  std::cout << "  --creatures     -c : The number of creatures in the population" << std::endl;
  std::cout << "  --steps         -s : The number of evolutionary ittereations" << std::endl;
  std::cout << "  --percent-mate  -pm: The percent of creatures to evolve to the next iteration" << std::endl;
  std::cout << "  --percent-kill  -pk: The percent of creatures to kill after an iteration" << std::endl;
  std::cout << "  --mutation-rate -mr: The std. dist. in the mutation generating function" << std::endl;
  std::cout << std::endl;
  std::cout << "                  -R : Reward for mutual cooperation (default 3)" << std::endl;
  std::cout << "                  -S : Sucker's payoff (default 0)" << std::endl;
  std::cout << "                  -T : Temptation to defect (default 5)" << std::endl;
  std::cout << "                  -P : Punishment for mutual defection (default 1)" << std::endl;
  std::cout << "                   Player B" << std::endl;
  std::cout << "                    C   D  " << std::endl;
  std::cout << "                  +-------+" << std::endl;
  std::cout << "                C | R | S |" << std::endl;
  std::cout << "       Player A   +---+---+" << std::endl;
  std::cout << "                D | T | P |" << std::endl; 
  std::cout << "                  +-------+" << std::endl;
}

void print_creature_weights(const thrust::device_vector<float> &creature_weights,
                            const int num_weights,
                            const int memory_length,
                            const bool use_const,
                            const bool use_random,
                            const bool use_same,
                            const unsigned int idx)
{
  std::cout << "  weights: ";
  thrust::copy(creature_weights.begin()+idx*num_weights,
               creature_weights.begin()+idx*num_weights+memory_length,
               std::ostream_iterator<float>(std::cout, ", "));
  unsigned int offset = idx*num_weights+memory_length;
  if(use_const)
  {
    std::cout << " const: " << creature_weights[offset];
    offset++;
  }
  if(use_random)
  {
    std::cout << " random: " << creature_weights[offset];
    offset++;
  }
  if(use_same)
  {
    std::cout << " same: " << creature_weights[offset];
    offset++;
  }
  std::cout << std::endl;
}

int main(int argc, char* argv[])
{
  event_pair timer;
  float cpu_time,gpu_time;
  ConfigFile *config;

  // Simulation Parameter Declarations
  bool verbose;
  bool run_on_cpu;
  int num_rounds_per_match;
  int memory_length;
  int num_weights;
  int num_creatures;
  int num_steps;
  float percent_mate;
  float percent_kill;
  float mutation_rate;
  //             Player B
  //              C   D
  //            +-------+
  //          C | R | S |
  // Player A   +---+---+
  //          D | T | P |
  //            +-------+
  int R; // Reward for mutual cooperation
  int S; // Sucker's payoff
  int T; // Temptation to defect
  int P; // Punishment for mutual defection

  bool use_const;
  bool use_random;
  bool use_same;

  char *output_config_filename = "saved_config.txt";
  char *input_config_filename = NULL;

  // Determine name of config file
  for (int i=1; i < argc; i++)
  {
    if (strcmp(argv[i], "--config-in") == 0 || strcmp(argv[i], "-ci") == 0)
    {
      input_config_filename = (++i < argc) ? argv[i] : NULL;
      break;
    }
  }


  // initialize the config object
  if(input_config_filename != NULL)
  {
    try
    {
      config = new ConfigFile(input_config_filename);
    } catch (...) {
      std::cerr << "Configuration file could not be found" << std::endl;
      config = new ConfigFile();
    }
  }
  else
  {
    config = new ConfigFile();
  }


  // read values from config, adding any defaults that were't in the file
  if(!config->readInto(verbose, "verbose", false))
    config->add("verbose", verbose);
  if(!config->readInto(run_on_cpu, "run_on_cpu", false))
    config->add("run_on_cpu", run_on_cpu);
  if(!config->readInto(num_rounds_per_match, "num_rounds_per_match", 101))
    config->add("num_rounds_per_match", num_rounds_per_match);
  if(!config->readInto(memory_length, "memory_length", 3))
    config->add("memory_length", memory_length);
  if(!config->readInto(num_creatures, "num_creatures", 512))
    config->add("num_creatures", num_creatures);
  if(!config->readInto(num_steps, "num_steps", 5))
    config->add("num_steps", num_steps);
  if(!config->readInto(percent_mate, "percent_mate", (float) 0.1))
    config->add("percent_mate", percent_mate);
  if(!config->readInto(percent_kill, "percent_kill", (float) 0.2))
    config->add("percent_kill", percent_kill);
  if(!config->readInto(mutation_rate, "mutation_rate", (float) 0.9))
    config->add("mutation_rate", mutation_rate);
  if(!config->readInto(R, "R", 3))
    config->add("R", R);
  if(!config->readInto(S, "S", 0))
    config->add("S", S);
  if(!config->readInto(T, "T", 5))
    config->add("T", T);
  if(!config->readInto(P, "P", 1))
    config->add("P", P);
  if(!config->readInto(use_const, "use_const", true))
    config->add("use_const", use_const);
  if(!config->readInto(use_random, "use_random", false))
    config->add("use_random", use_random);
  if(!config->readInto(use_same, "use_same", false))
    config->add("use_same", use_same);

  // read command-line overrides, replacing existing config values
  for (int i=1; i < argc; i++)
  {
    if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0)
    {
      print_help();
      exit(0);
    }
    else if (strcmp(argv[i], "--version") == 0)
    {
      print_version();
      exit(0);
    }
    else if (strcmp(argv[i], "--config-in") == 0 || strcmp(argv[i], "-ci") == 0)
    {
      i++;
      continue;
    }
    else if (strcmp(argv[i], "--config-out") == 0 || strcmp(argv[i], "-co") == 0)
    {
      output_config_filename = (++i < argc) ? argv[i] : output_config_filename;
      std::cerr << "Specified --config-out or -co but no value. Using " << output_config_filename << " instead." << std:: endl;
    }
    else if (strcmp(argv[i], "--verbose") == 0 || strcmp(argv[i], "-v") == 0)
    {
      verbose = true;
      config->add("verbose", verbose);
    }
    else if (strcmp(argv[i], "--cpu") == 0)
    {
      run_on_cpu = true;
      config->add("run_on_cpu", run_on_cpu);
    }
    else if (strcmp(argv[i], "--rounds") == 0 || strcmp(argv[i], "-r") == 0)
    {
      num_rounds_per_match = (++i < argc) ? atoi(argv[i]) : 0;
      config->add("num_rounds_per_match", num_rounds_per_match);
    }
    else if (strcmp(argv[i], "--memory") == 0 || strcmp(argv[i], "-m") == 0)
    {
      memory_length = (++i < argc) ? atoi(argv[i]) : 0;
      config->add("memory_length", memory_length);
    }
    else if (strcmp(argv[i], "--creatures") == 0 || strcmp(argv[i], "-c") == 0)
    {
      num_creatures = (++i < argc) ? atoi(argv[i]) : 0;
      config->add("num_creatures", num_creatures);
    }
    else if (strcmp(argv[i], "--steps") == 0 || strcmp(argv[i], "-s") == 0)
    {
      num_steps = (++i < argc) ? atoi(argv[i]) : 0;
      config->add("num_steps", num_steps);
    }
    else if (strcmp(argv[i], "--percent-mate") == 0 || strcmp(argv[i], "-pm") == 0)
    {
      percent_mate = (++i < argc) ? atof(argv[i]) : 0;
      config->add("percent_mate", percent_mate);
    }
    else if (strcmp(argv[i], "--percent-kill") == 0 || strcmp(argv[i], "-pk") == 0)
    {
      percent_kill = (++i < argc) ? atof(argv[i]) : 0;
      config->add("percent_kill", percent_kill);
    }
    else if (strcmp(argv[i], "--mutation-rate") == 0 || strcmp(argv[i], "-mr") == 0)
    {
      mutation_rate = (++i < argc) ? atof(argv[i]) : 0;
      config->add("mutation_rate", mutation_rate);
    }
    else if (strcmp(argv[i], "-R") == 0)
    {
      R = (++i < argc) ? atoi(argv[i]) : 0;
      config->add("R",R);
    }
    else if (strcmp(argv[i], "-S") == 0)
    {
      S = (++i < argc) ? atoi(argv[i]) : 0;
      config->add("S",S);
    }
    else if (strcmp(argv[i], "-T") == 0)
    {
      T = (++i < argc) ? atoi(argv[i]) : 0;
      config->add("T",T);
    }
    else if (strcmp(argv[i], "-P") == 0)
    {
      P = (++i < argc) ? atoi(argv[i]) : 0;
      config->add("P",P);
    }
    else if (strcmp(argv[i], "--use_const") == 0)
    {
      use_const = true;
      config->add("use_const",use_const);
    }
    else if (strcmp(argv[i], "--no_use_const") == 0)
    {
      use_const = false;
      config->add("use_const",use_const);
    }
    else if (strcmp(argv[i], "--use_random") == 0)
    {
      use_random = true;
      config->add("use_random",use_random);
    }
    else if (strcmp(argv[i], "--no_use_random") == 0)
    {
      use_random = false;
      config->add("use_random",use_random);
    }
    else if (strcmp(argv[i], "--use_same") == 0)
    {
      use_same = true;
      config->add("use_same",use_same);
    }
    else if (strcmp(argv[i], "--no_use_same") == 0)
    {
      use_same = false;
      config->add("use_same",use_same);
    }
    else
    {
      std::cerr << "Unknown parameter: " << argv[i] << std::endl;
      print_help();
      exit(-1);
    }
  }


  // Validate everything
  if(!(num_rounds_per_match > 0))
  {
    std::cerr << "Number of rounds has to be positive." << std::endl;
    exit(-1);
  }
  if(!(memory_length >= 0 && memory_length <= 10))
  {
    std::cerr << "Memory has to be between 0 and 10." << std::endl;
    exit(-1);
  }
  if(!(num_creatures > 0))
  {
    std::cerr << "Number of creatures has to be positive." << std::endl;
    exit(-1);
  }
  if(!(num_steps > 0))
  {
    std::cerr << "Number of evolution steps has to be positive." << std::endl;
    exit(-1);
  }
  if(!(percent_mate >= 0.0f && percent_mate <= 1.0f))
  {
    std::cerr << "percent-mate has to be between 0 and 1." << std::endl;
    exit(-1);
  }
  if(!(percent_kill >= 0.0f && percent_kill <= 1.0f))
  {
    std::cerr << "percent-kill has to be between 0 and 1." << std::endl;
    exit(-1);
  }
  if(!(mutation_rate >= 0.0f && mutation_rate <= 10.0f))
  {
    std::cerr << "mutation-rate has to be between 0 and 10.0" << std::endl;
    exit(-1);
  }


  // Write the configuration back out to a new file, the name of which
  // can be set as a parameter or is otherwise "saved_config.txt"
  std::ofstream *outputFileStream = new std::ofstream(output_config_filename, std::ofstream::trunc);
  if(!outputFileStream->is_open())
  {
    std::cerr << "could not open output parameter file named " << output_config_filename << std::endl;
    exit(-1);
  }
  *outputFileStream << *config;
  outputFileStream->close();
  delete outputFileStream;


  // done with the config object, so we can delete it
  delete config;


  // Tidy up the initialization
  num_weights = memory_length;
  if(use_const)  num_weights++;  // Add any extra weights as needed
  if(use_random) num_weights++;
  if(use_same)   num_weights++;

  thrust::host_vector<float> h_creature_weights(num_creatures * num_weights);
  thrust::host_vector<int> h_creature_scores(num_creatures);
  thrust::host_vector<unsigned int> h_random(num_creatures * num_creatures);

  thrust::device_vector<float> d_creature_weights(num_creatures * num_weights);
  thrust::device_vector<int> d_creature_scores(num_creatures);
  thrust::device_vector<unsigned int> d_random(num_creatures * num_creatures);

  if(use_random) {
    if(run_on_cpu)
      thrust::copy(thrust::make_counting_iterator(0), thrust::make_counting_iterator(num_creatures * num_creatures), h_random.begin());
    
      thrust::copy(thrust::make_counting_iterator(0), thrust::make_counting_iterator(num_creatures * num_creatures), d_random.begin());
  }

  // convert percents into numbers of elements
  unsigned int num_mate = (unsigned int) (percent_mate * num_creatures);
  unsigned int num_kill = (unsigned int) (percent_kill * num_creatures);

  // generate random weights
  thrust::generate(h_creature_weights.begin(), h_creature_weights.end(), random_float);
  thrust::fill(h_creature_weights.begin(), h_creature_weights.begin()+3*num_weights, (float) 0);
  h_creature_weights[0*num_weights+num_weights-1] = 1.0;
  h_creature_weights[1*num_weights+num_weights-1] = -1.0;
  h_creature_weights[2*num_weights] = 1.0;
  h_creature_weights[2*num_weights+num_weights-1] = 0.1;

//PRINf(h_creature_weights);

  // initialize device vectors
  d_creature_weights = h_creature_weights;

  // # evolutionary steps
  for(unsigned int step = 1; step <= num_steps; step++) {
    std::cout << std::endl << "Evolutionary step #" << step << std::endl;
    
    // normalize weights per creature on [-1.0,1.0]
  /*
    normalize_all_weights(num_creatures,
                          num_weights,
                          d_creature_weights,
                          h_creature_weights);
    */
    if (run_on_cpu)
    {
      // reset scores
      thrust::fill(h_creature_scores.begin(), h_creature_scores.end(), (int) 0);
    
      // run sim on the host
      start_timer(&timer);
      
      host_simulation(memory_length,
                      num_weights,
                      num_creatures,
                      num_rounds_per_match,
                      h_random,
                      h_creature_weights,
                      h_creature_scores,
                      R,
                      S,
                      T,
                      P,
                      use_const,
                      use_random,
                      use_same);
      
      cpu_time = stop_timer(&timer, "CPU");
    }
    
    // run sim on device
    start_timer(&timer);
    device_simulation(memory_length,
                      num_weights,
                      num_creatures,
                      num_rounds_per_match,
                      d_random,
                      d_creature_weights,
                      d_creature_scores,
                      R,
                      S,
                      T,
                      P,
                      use_const,
                      use_random,
                      use_same);
    gpu_time = stop_timer(&timer, "GPU");
  
    if(run_on_cpu)
    {
      if(h_creature_scores != d_creature_scores) {
        std::cout << "assert failed!" << std::endl;
        PRINT(h_creature_scores);
        PRINT(d_creature_scores);
        exit(-1);
      }

      std::cout << " Speedup: " << cpu_time / gpu_time << "x" << std::endl;
    }

    unsigned int max_idx = thrust::max_element(d_creature_scores.begin(), d_creature_scores.end()) - d_creature_scores.begin();
    std::cout << "  Fittest: creature number " << max_idx << std::endl;
    std::cout << "    score: " << d_creature_scores[max_idx] << std::endl;
    print_creature_weights(d_creature_weights, num_weights, memory_length, use_const, use_random, use_same, max_idx);
    //PRIN2f(d_creature_weights.begin()+max_idx*num_weights, d_creature_weights.begin()+(max_idx+1)*num_weights);
          
    unsigned int min_idx = thrust::min_element(d_creature_scores.begin(), d_creature_scores.end()) - d_creature_scores.begin();
    std::cout << "Sickliest: creature number " << min_idx << std::endl;
    std::cout << "    score: " << d_creature_scores[min_idx] << std::endl;
    print_creature_weights(d_creature_weights, num_weights, memory_length, use_const, use_random, use_same, min_idx);

//    PRINT(d_creature_scores);

//    PRINT(h_random);

    // update creatures, mating the fittest, killing off the weakest
    // 1) sort the creatures by their scores
    if(!run_on_cpu) {
        h_creature_scores = d_creature_scores;
        // the weights are copied after the normalization step, so it's unnecessary here
    }

    thrust::host_vector<int> h_creature_indices(num_creatures);
    thrust::copy(thrust::make_counting_iterator(0), thrust::make_counting_iterator(num_creatures), h_creature_indices.begin());
    thrust::sort_by_key(h_creature_scores.begin(), h_creature_scores.end(), h_creature_indices.begin());

    if(verbose)
    {
      std::cout << "--------" << std::endl;
      for (int i=num_creatures-1; 0<=i; i--) {
          unsigned int t = h_creature_indices[i];
          std::cout << "[" << t << ", " << h_creature_scores[i] << "]";
          print_creature_weights(d_creature_weights, num_weights, memory_length, use_const, use_random, use_same, t);
      }
    }

    perturb_with_normal perturbation_op(mutation_rate);

    for(int i=0; i<num_kill; i++)
    {
      // 2) use the sorted indices to update the last num_kill creatures with copies of the first num_mate creatures
      int m = h_creature_indices[num_creatures - (i % num_mate) - 1];
      int k = h_creature_indices[i];
      thrust::copy(h_creature_weights.begin() + m * num_weights,
                   h_creature_weights.begin() + (m+1) * num_weights,
                   h_creature_weights.begin() + k * num_weights);

      // 3) evolve all of the parameters
      thrust::transform(h_creature_weights.begin() + k * num_weights,
                        h_creature_weights.begin() + (k+1) * num_weights,
                        h_creature_weights.begin() + k * num_weights,
                        perturbation_op);
    }

    // 4) copy new weights from host to device
    d_creature_weights = h_creature_weights;
  }
  
  return 0;
}
